import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from environment import GridWorld
from collections import Counter


class DQN(nn.Module):
    """定义q网络模型:
    s-->q0, q1, q2, q3, q4"""

    def __init__(self, input_size, hidden_size, output_size):
        super(DQN, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.nonlinear1 = nn.Sigmoid()
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.nonlinear2 = nn.Sigmoid()
        self.fc3 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = self.fc1(x)
        x = self.nonlinear1(x)
        x = self.fc2(x)
        x = self.nonlinear2(x)
        x = self.fc3(x)
        return x


class ReplayBuffer:
    """定义经验回放缓存"""

    def __init__(self, init_size, capacity):
        self.capacity = capacity
        self.buffer = []
        self.buffer_init(GridWorld(5), init_size)

    def push(self, transition):
        self.buffer.append(transition)
        if len(self.buffer) > self.capacity:
            self.buffer.pop(0)

    def buffer_init(self, env: GridWorld, num: int = None):
        """初始填入一定量的数据"""
        num = 0 if num is None else num
        for _ in range(num):
            idx = np.random.choice(env.size**2)
            s = env.state_space[idx]
            a = np.random.choice(env.action_space)
            r = env._reward(s, a, env.size)
            next_s = env._next_state(s, a, env.size)
            self.push((s, a, r, next_s))

    def sample(self, batch_size):
        choosed_idxs = np.random.choice(range(len(self.buffer)), batch_size, replace=False)
        batch = [self.buffer[i] for i in choosed_idxs]
        states, actions, rewards, next_states = zip(*batch)
        return (
            np.array(states),
            np.array(actions),
            np.array(rewards),
            np.array(next_states),
        )

    # 统计sars'分布
    def counter_sa(self):
        counter = Counter()
        idxs = []
        for exp in self.buffer:
            idxs.append((tuple(exp[0]), exp[1]))
        counter.update(idxs)
        return counter


# 定义DQN算法
class DQNAgent:
    def __init__(
        self,
        state_size,
        action_size,
        hidden_size,
        lr,
        lr_milestones,
        gamma,
        buffer_init_size,
        buffer_capacity,
        batch_size,
    ):
        self.state_size = state_size
        self.action_size = action_size
        self.hidden_size = hidden_size
        self.lr = lr
        self.lr_miletones = lr_milestones

        self.gamma = gamma
        self.batch_size = batch_size
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self.model = DQN(state_size, hidden_size, action_size).to(self.device)
        self.target_model = DQN(state_size, hidden_size, action_size).to(self.device)
        self.target_model.load_state_dict(self.model.state_dict())
        self.target_model.eval()

        self.optimizer = optim.Adam(self.model.parameters(), lr=self.lr)
        self.scheduler = torch.optim.lr_scheduler.MultiStepLR(
            self.optimizer,
            milestones=self.lr_miletones,
            gamma=0.8,
            last_epoch=-1,
        )
        self.loss_fn = nn.MSELoss()
        self.replay_buffer = ReplayBuffer(init_size=buffer_init_size, capacity=buffer_capacity)
        # 初始化经验池

    def update_target_model(self):
        self.target_model.load_state_dict(self.model.state_dict())

    def train(self):
        """训练一个batch, 返回该batch的TD误差"""
        if len(self.replay_buffer.buffer) < self.batch_size:
            return

        states, actions, rewards, next_states = self.replay_buffer.sample(self.batch_size)

        states = torch.tensor(states).float().to(self.device)
        actions = torch.tensor(actions).long().to(self.device)
        rewards = torch.tensor(rewards).float().to(self.device)
        next_states = torch.tensor(next_states).float().to(self.device)
        q_values = self.model(states).gather(1, actions.unsqueeze(1))
        with torch.no_grad():
            next_q_values = self.target_model(next_states).max(1)[0]
        # breakpoint()
        expected_q_values = rewards + self.gamma * next_q_values
        # breakpoint()
        loss = self.loss_fn(q_values, expected_q_values.unsqueeze(1))
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        return loss.item()

    def get_lr(self):
        '''返回实时学习率'''
        return self.optimizer.param_groups[0]["lr"]

    def get_action(self, state, epsilon):
        """策略在线时，用于选择动作，ε-greedy策略"""
        if np.random.rand() <= epsilon:
            return np.random.randint(self.action_size)
        else:
            state = torch.tensor(state).float().unsqueeze(0).to(self.device)
            with torch.no_grad():
                q_values = self.model(state)
            return torch.argmax(q_values, dim=1).item()
